# 多种群遗传算法
import random
import numpy as np
import csv
import time


class GAMp:
    def __init__(self, fitness_func, n, num_populations, population_size, generations, casei):
        self.n = n  # 编码长度
        self.num_populations = num_populations
        self.population_size = population_size
        self.generations = generations
        self.fitness_func = fitness_func  # 将fitness作为函数参数
        self.casei = casei

        self.populations = []
        self.y_populations = []
        self.best_fitness = -np.inf  # 初始化最佳适应度为无穷小量
        self.number = 0  # 初始化迭代次数

        for _ in range(self.num_populations):
            population = np.zeros((self.population_size, self.n))
            y_population = np.zeros(self.population_size)
            self.populations.append(population)
            self.y_populations.append(y_population)
            self.initialize_population(population)

    def initialize_population(self, population):
        population[:] = np.random.randint(0, 2, size=(self.population_size, self.n))

    def cal_y(self):
        for i in range(self.num_populations):
            self.y_populations[i] = np.array([self.fitness_func(self.casei, x) for x in self.populations[i]])

        current_max = np.max([np.max(y) for y in self.y_populations])
        yy = np.array(self.y_populations)
        if current_max > self.best_fitness:
            self.best_fitness = current_max
            max_indices = np.unravel_index(np.argmax(yy), yy.shape)
            aa, bb = int(max_indices[0]), int(max_indices[1])
            self.best_code = self.populations[aa][bb]

            BestCode_result_file = "GAMP_code_" + str(self.casei) + ".csv"
            with open(BestCode_result_file, 'w', encoding='utf-8', newline='') as f:
                write = csv.writer(f)
                write.writerow(list(self.best_code))

    def selection(self):
        selected_populations = []
        for i in range(self.num_populations):
            # 使用numpy生成竞争者索引矩阵
            competitors = np.random.randint(0, self.population_size, (self.population_size, 3))
            # 获取每组竞争者的适应度
            competitors_fitness = self.y_populations[i][competitors]
            # 获取获胜者的索引
            winners = competitors[np.arange(len(competitors)), np.argmax(competitors_fitness, axis=1)]
            # 选择获胜者
            selected_populations.append(self.populations[i][winners])
        self.populations = selected_populations

    def crossover(self):
        for num in range(self.num_populations):
            # 生成交叉点
            crossover_points = np.random.randint(1, self.n - 1, self.population_size // 2)
            # 重排population以便配对
            pop = self.populations[num]
            indices = np.arange(self.population_size)
            np.random.shuffle(indices)
            pop = pop[indices]

            # 将种群分成两半进行配对
            half = self.population_size // 2
            parents1 = pop[:half]
            parents2 = pop[half:2 * half]

            # 创建子代
            children = np.zeros((self.population_size, self.n))
            for i, cp in enumerate(crossover_points):
                children[2 * i] = np.concatenate((parents1[i, :cp], parents2[i, cp:]))
                children[2 * i + 1] = np.concatenate((parents2[i, :cp], parents1[i, cp:]))

            self.populations[num] = children

    def mutation(self):
        mutation_rate = 0.001
        # 生成布尔类型的mutation_mask
        mutation_mask = np.random.random(
            (self.num_populations, self.population_size, self.n)) < mutation_rate
        # 确保populations和mutation_mask的数据类型一致
        self.populations = [np.logical_xor(pop.astype(bool), mutation_mask[i]).astype(int)
                            for i, pop in enumerate(self.populations)]

    def migration(self):
        # 迁移操作，将一部分个体从一个种群迁移到另一个种群
        for i in range(self.num_populations):
            population = self.populations[i]
            migration_indices = random.sample(range(self.population_size), int(self.population_size * 0.1))
            for j in migration_indices:
                target_population_index = (i + 1) % self.num_populations
                target_population = self.populations[target_population_index]
                target_population[random.randint(0, self.population_size - 1)] = population[j]

    def run(self):
        for _ in range(self.generations):
            sti = time.time()
            self.number += 1  # 迭代次数加一
            self.cal_y()  # 计算当前种群的fitness
            self.selection()
            self.crossover()
            self.mutation()
            self.migration()
            Curve_result_file = "GAMP_result_" + str(self.casei) + ".csv"
            with open(Curve_result_file, 'a', encoding='utf-8', newline='') as f:
                write = csv.writer(f)
                write.writerow([self.number * self.population_size * self.num_populations, self.best_fitness.round(5)])

            print("Iter ", _, ": ", self.best_fitness, "Running time: ", time.time() - sti, 's')
        population = self.populations[-1]
        y_population = self.y_populations[-1]
        max_indices = np.unravel_index(np.argmax(y_population), y_population.shape)
        best_code = population[tuple(max_indices)]  # 输出适应度最大对应的最优编码
        np.set_printoptions(threshold=np.inf)  # 显示数组的全部元素
        # with open("GAMp_result.csv", 'a', encoding='utf-8', newline='') as f:  # 打开文件并记录评估次数和最佳适应度
        #     write = csv.writer(f)  # 创建writer对象
        #     write.writerow(list([self.best_fitness, best_code]))  # 记录最佳适应度和基因序列
        return best_code, self.best_fitness  # 返回最佳适应度和基因序列

